import os
import json
import logging
import time
import base64
from openai import OpenAI
import pandas as pd
from config import settings
from PIL import Image


# 配置 OpenAI 代理 API
client = OpenAI(
    base_url=settings.base_url,
    api_key=settings.api_key,
)

root_dir = settings.root_dir
# 输入文件
INPUT_JSON = f"{root_dir}/dataset/forgery/forgery.json"

# 配置日志
LOG_FILE = f"{root_dir}/log/patch_generate.log"
os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)
logging.basicConfig(filename=LOG_FILE, level=logging.INFO, 
                    format='%(asctime)s - %(levelname)s - %(message)s')
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
logging.getLogger().addHandler(console_handler)

# 输出 JSON 文件
BASE_DIR = f"{root_dir}/qa_design"
OUTPUT_JSON = f"{BASE_DIR}/patch_generate.json"

# prompt_file
PROMPT_FILE = f"{BASE_DIR}/patch_generate_prompt.txt"
with open(PROMPT_FILE, "r", encoding="utf-8") as f:
    prompt_template = f.read()

def load_existing_results():
    """ 加载已生成的结果，支持断点续传 """
    if os.path.exists(OUTPUT_JSON):
        with open(OUTPUT_JSON, "r", encoding="utf-8") as f:
            try:
                return json.load(f)
            except json.JSONDecodeError:
                return []
    return []

def save_results(data):
    """ 保存当前结果 """
    with open(OUTPUT_JSON, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)

def compress_image(image_path, quality=85):
    """压缩图片"""
    image = Image.open(image_path)
    image.save(image_path, optimize=True, quality=quality)
    return image_path

def encode_image(image_path):
    # 
    """将图片编码为base64格式"""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def get_answer(image_path):
    """使用GPT-4o-mini生成图片描述"""
    logging.info(f"处理图片: {image_path}")
    image_path = os.path.join(root_dir, image_path)
    panel_path = os.path.join(root_dir, image_path.replace("mask", "panel"))
    if not os.path.exists(image_path):
        logging.error(f"掩码不存在: {image_path}")
        return None

    if not os.path.exists(panel_path):
        logging.error(f"panel不存在: {panel_path}")
        return None

    # 需要根据panel大小，将图片resize到panel大小
    panel_image = Image.open(panel_path)
    panel_width, panel_height = panel_image.size
    image_image = Image.open(image_path)
    image_image = image_image.resize((panel_width, panel_height))
    
    try:
        quality = 85
        # 如果图片大于25MB，则压缩图片，直到小于25MB
        while os.path.getsize(image_path) >= 20 * 1024 * 1024:
            quality -= 5
            image_path = compress_image(image_path, quality=quality)
            logging.info(f"压缩图片: {image_path}")
            time.sleep(1)   

        # 将图片编码为base64
        base64_image = encode_image(image_path)
        
        # 构建提示词
        prompt = prompt_template

        # 调用API
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}"
                            }
                        }
                    ]
                }
            ],
            max_tokens=800
        )
        
        description = response.choices[0].message.content
        logging.info(f"成功生成描述")
        return description
    except Exception as e:
        logging.error(f"生成描述失败: {str(e)}")
        return None


def process_image_list(panels):
    """处理面板列表，生成描述并保存结果"""
    existing_results = load_existing_results()
    existing_panel_ids = {item["panel_id"] for item in existing_results}
    
    total_panels = len(panels)
    results = existing_results.copy()
    written_image_paths = []

    for index, panel in enumerate(panels, start=1):
        panel_id = panel["panel_id"]
        image_path = os.path.join(root_dir, panel["mask_path"])

        # 检查是否已处理过
        if panel_id in existing_panel_ids:
            logging.info(f"跳过 {panel_id}, 已处理.")
            continue
        
        logging.info(f"处理 {index}/{total_panels}: {panel_id}")
        
        # 获取图片描述
        answer = get_answer(image_path)
        
        if not answer:
            logging.error(f"无法获取描述: {panel_id}")
            continue
        
        # 添加到结果中
        result = {
            "panel_id": panel_id,
            "mask_path": panel["mask_path"],
            "answer": answer
        }
        
        results.append(result)
        save_results(results)  # 实时保存进度

        logging.info(f"完成 {index}/{total_panels}: {panel_id}")
        time.sleep(2)  # 避免触发 API 速率限制
    
    return results


def load_data(target_issue_id="FG_004"):
    """
    从 JSON 中提取 panel_id + mask_path
    条件：
      - panel_id 包含 target_issue_id
      - panel_level_issues.has_issue == True
    """
    with open(INPUT_JSON, "r", encoding="utf-8") as f:
        data = json.load(f)

    results = []

    for paper in data.get("papers", []):
        for page in paper.get("pages", []):
            for fig in page.get("figures", []):
                for panel in fig.get("panels", []):
                    panel_id = panel.get("panel_id", "")
                    panel_issues = panel.get("panel_level_issues", {})

                    if target_issue_id in panel_id and panel_issues.get("has_issue", False):
                        for issue in panel_issues.get("issues", []):
                            evidence = issue.get("evidence", {})
                            mask_path = evidence.get("mask_path")
                            if mask_path:
                                results.append({
                                    "panel_id": panel_id,
                                    "mask_path": mask_path
                                })

    return results

if __name__ == "__main__":
    data = load_data("FG_004")
    process_image_list(data[:40])